#include "kernel_operator.h"
using namespace AscendC;
constexpr int32_t BUFFER_NUM = 2;                                     // tensor num for each queue

template<typename typeT>
class KernelAsinh {
public:
    __aicore__ inline KernelAsinh() {}
    __aicore__ inline void Init(GM_ADDR x, GM_ADDR y, uint32_t inputNum, 
                                uint32_t tileLength) {
        ASSERT(GetBlockNum() != 0 && "block dim can not be zero!");

        this->inputNum = inputNum;
        this->tileLength = tileLength;

        xGm.SetGlobalBuffer((__gm__ typeT*)x, this->inputNum + 16);
        yGm.SetGlobalBuffer((__gm__ typeT*)y, this->inputNum + 16);

        // pipe alloc memory to queue, the unit is Bytes
        pipe.InitBuffer(inQueueX, BUFFER_NUM, this->tileLength * sizeof(typeT));
        pipe.InitBuffer(outQueueY, BUFFER_NUM, this->tileLength * sizeof(typeT));
        pipe.InitBuffer(tmp1, this->tileLength * sizeof(float));
        pipe.InitBuffer(tmp2, this->tileLength * sizeof(float));
        pipe.InitBuffer(tmp3, this->tileLength * sizeof(int8_t) / 8);
    }
    __aicore__ inline void Process() {
        int32_t loopCount = this->inputNum / this->tileLength;
        if(loopCount == 0) {
            uint32_t length = (this->inputNum + 15) / 16 * 16;
            CopyIn(0, length);
            Compute(0, this->inputNum);
            CopyOut(0, length);
        } else {
            for (int32_t i = 0; i < loopCount; i++) {
                CopyIn(i, this->tileLength);
                Compute(i, this->tileLength);
                CopyOut(i, this->tileLength);
            } 
            if (this->inputNum % this->tileLength > 0) {
                uint32_t length = this->inputNum % this->tileLength;
                uint32_t lengthalign = (length + 15) / 16 * 16;
                CopyIn(loopCount, lengthalign);
                Compute(loopCount, length);
                CopyOut(loopCount, lengthalign);
            }   
        }        
    }

private:
    __aicore__ inline void CopyIn(int32_t progress, uint32_t length) {
        LocalTensor<typeT> xLocal = inQueueX.AllocTensor<typeT>();
        DataCopy(xLocal, xGm[progress * this->tileLength], length);
        inQueueX.EnQue(xLocal);
    }
    __aicore__ inline void Compute(int32_t progress, uint32_t length) {
        LocalTensor<typeT> xLocal = inQueueX.DeQue<typeT>();
        LocalTensor<typeT> yLocal = outQueueY.AllocTensor<typeT>();

        if constexpr (std::is_same_v<typeT, float>) {

            // auto p1 = tmp1.Get<float>();
            // Abs(p1, xLocal, length);
            // Div(p1, p1, xLocal, length);
            // Abs(xLocal, xLocal, length);
            LocalTensor<typeT> tmpLT1 = tmp1.Get<typeT>();
            LocalTensor<int8_t> cmp1 = tmp3.Get<int8_t>();
            Duplicate(tmpLT1, typeT(0), length);
            Compare(cmp1, xLocal, tmpLT1, CMPMODE::GE, length);  
            Duplicate(tmpLT1, typeT(1), length);      
            Select(tmpLT1, cmp1, tmpLT1, typeT(-1), SELMODE::VSEL_TENSOR_SCALAR_MODE, length);
            Mul(xLocal, xLocal, tmpLT1, length);

            Mul(yLocal, xLocal, xLocal, length);
            Adds(yLocal, yLocal, typeT(1), length);
            Sqrt(yLocal, yLocal, length);
            Add(yLocal, yLocal, xLocal, length);
            Ln(yLocal, yLocal, length);

            Mul(yLocal, yLocal, tmpLT1, length);
 
        } else if constexpr (std::is_same_v<typeT, half>) {
            auto p1 = tmp1.Get<float>();
            auto p2 = tmp2.Get<float>();
            Cast(p1, xLocal, RoundMode::CAST_NONE, length);
            Mul(p2, p1, p1, length);
            Adds(p2, p2, float(1), length);
            Sqrt(p2, p2, length);
            Add(p2, p2, p1, length);
            Ln(p2, p2, length);
            Cast(yLocal, p2, RoundMode::CAST_ROUND, length);
        }

        outQueueY.EnQue<typeT>(yLocal);
        inQueueX.FreeTensor(xLocal);
    }
    __aicore__ inline void CopyOut(int32_t progress, uint32_t length) {
        LocalTensor<typeT> yLocal = outQueueY.DeQue<typeT>();
        DataCopy(yGm[progress * this->tileLength], yLocal, length);
        outQueueY.FreeTensor(yLocal);
    }

private:
    TPipe pipe;
    TBuf<QuePosition::VECCALC> tmp1, tmp2, tmp3;
    TQue<QuePosition::VECIN, BUFFER_NUM> inQueueX;
    TQue<QuePosition::VECOUT, BUFFER_NUM> outQueueY;
    GlobalTensor<typeT> xGm;
    GlobalTensor<typeT> yGm;
    uint32_t inputNum;
    uint32_t tileLength;

};

extern "C" __global__ __aicore__ void asinh(GM_ADDR x, GM_ADDR y, GM_ADDR workspace, GM_ADDR tiling) {
    GET_TILING_DATA(tiling_data, tiling);
    KernelAsinh<DTYPE_X> op;
    op.Init(x, y, tiling_data.inputNum, tiling_data.tileLength);
    op.Process();
}